#include "Rsa.h"
#include <cstring>
#include <cmath>
#include <iostream>
#include "BinaryIO.h"

namespace Utils::Rsa
{
    typedef union
    {
        char buff[8];
        struct
        {
            int32_t n;
            int32_t key;
        };
        int64_t value;
    } R8DATA;

    static int32_t primer[6541];

    static int32_t IsPrimer(int32_t n)
    {
        int32_t m = (int32_t)sqrt(n);
        for (int32_t index = 0; primer[index] <= m; index++)
        {
            if (n % primer[index] == 0)
            {
                return false;
            }
        }
        return true;
    }

    static int32_t Gcd(int32_t a, int32_t b)
    {
        for (int32_t t; b; t = a % b, a = b, b = t)
            ;
        return a;
    }

    static void ExGcd(int32_t a, int32_t b, int32_t &x, int32_t &y)
    {
        if (b == 0) // 当b为0时,这也是欧几里得算法终止的条件,此时x1=1,y1=0
        {
            x = 1, y = 0;
            return;
        }
        ExGcd(b, a % b, x, y);    // 没有终止,再向下再次递归,参数改为b,a%b
        int32_t t = x;            // 已经计算出了b*x1+(a%b)*y1=gcd(b,a)
        x = y, y = t - a / b * y; // 可以推导出x,y
    }

    static void GetPrimers(int32_t &p, int32_t &q)
    {
        srand(time(NULL));
        do
        {
            p = rand() % 10000 + 30000;
        } while (IsPrimer(p) == false);

        do
        {
            q = rand() % 10000 + 30000;
        } while (IsPrimer(q) == false);
    }

    static int32_t Mod(uint64_t a, uint64_t b, uint64_t n)
    {
        uint64_t t = 1;
        while (b)
        {
            if (b % 2)
            {
                t = (t * a) % n;
            }
            a = (a * a) % n;
            b /= 2;
        }
        return (int32_t)t;
    }

    static void Init()
    {
        bool tag[65536];
        memset(tag, 0x00, sizeof(tag));
        int32_t i, j;
        for (i = 3; i < 256; i += 2)
        {
            if (tag[i] == false)
            {
                for (j = i * i; j < 65536; j += i)
                {
                    tag[j] = true;
                }
            }
        }     
        for (primer[0] = 2, j = 1, i = 3; i < 65536; i += 2)
        {
            if (tag[i] == false)
            {
                primer[j++] = i;
            }
        }
    }

    void CreateKey(int64_t &publicKey, int64_t &privateKey)
    {
        int32_t p, q, pub_key, pri_key, m, n, t; // pub_key公钥,pri_kry私钥
        Init();
        do
        {
            GetPrimers(p, q);
            n = p * q;
            m = (p - 1) * (q - 1);
            do
            {
                pub_key = (rand() % 1000 + 500) * (rand() % 1000 + 500);
            } while (Gcd(m, pub_key) != 1);

            ExGcd(pub_key, m, pri_key, t); // pub_key*pri_key=1mod(m);

        } while (pri_key < 1e5);

        R8DATA r8data;
        r8data.key = pub_key;
        r8data.n = n;
        publicKey = r8data.value;

        r8data.key = pri_key;
        r8data.n = n;
        privateKey = r8data.value;
    }

    uint32_t Encrypt(const char *inputData, uint32_t inputSize, int64_t key, char *outputBuff, uint32_t outputSize)
    {
        R8DATA r8data;
        r8data.value = key;

        uint16_t clearText[10000];
        int32_t cipherText[10000];

        uint32_t inputIndex = 0;
        uint32_t outputIndex = 0;
        uint32_t len;

        MemoryBinaryWriter writer(outputBuff, outputSize);
        writer.writeUInt32(inputSize);
        do
        {
            memset(clearText, 0, sizeof(clearText));
            if (inputSize - inputIndex >= 20000)
            {
                len = 20000;
            }
            else
            {
                len = inputSize - inputIndex;
            }
            memcpy(clearText, inputData, len);
            inputIndex += len;
            if (len % 2 != 0)
            {
                len++;
            }
            len /= 2;
            for (int32_t index = 0; index < len; ++index)
            {
                writer.writeInt32((int32_t)Mod((uint16_t)clearText[index], r8data.key, r8data.n));
            }
        } while (inputIndex < inputSize);

        return writer.getPosition();
    }

    uint32_t Decrypt(const char *inputData, uint32_t inputSize, int64_t key, char *outputBuff, uint32_t outputSize)
    {
        R8DATA r8data;
        r8data.value = key;

        MemoryBinaryReader reader(inputData, inputSize);
        MemoryBinaryWriter writer(outputBuff, outputSize);

        uint32_t total = reader.readUInt32();
        uint32_t len = total / 2;

        for (uint32_t index = 0; index < len; ++index)
        {
            int32_t cipherValue = reader.readInt32();
            writer.writeUInt16((uint16_t)Mod((uint32_t)cipherValue, r8data.key, r8data.n));
        }
        if (total % 2 != 0)
        {
            int32_t cipherValue = reader.readInt32();
            uint16_t value = (uint16_t)Mod((uint32_t)cipherValue, r8data.key, r8data.n);
            uint8_t value2 = (value >> 8) & 0xFF;
            writer.writeUInt8(value2);
        }

        return writer.getPosition();
    }

}